#include <OctomapGen.h>

using namespace octomap;
using octomap_msgs::OctomapBinary;

namespace octomap_gen {

OctomapGen::OctomapGen(ros::NodeHandle private_nh_)
: m_nh(),
  m_pointCloudSub(NULL),
  m_tfPointCloudSub(NULL),
  m_octree(NULL),
  m_maxRange(-1.0),
  m_worldFrameId("/map"),
  m_baseFrameId("base_footprint"),
  m_res(0.05),
  m_treeDepth(0),
  m_maxTreeDepth(0),
  m_probHit(0.7), m_probMiss(0.4),
  m_thresMin(0.12), m_thresMax(0.97),
  m_pointcloudMinX(-std::numeric_limits<double>::max()),
  m_pointcloudMaxX(std::numeric_limits<double>::max()),
  m_pointcloudMinY(-std::numeric_limits<double>::max()),
  m_pointcloudMaxY(std::numeric_limits<double>::max()),
  m_pointcloudMinZ(-std::numeric_limits<double>::max()),
  m_pointcloudMaxZ(std::numeric_limits<double>::max())
{
  ros::NodeHandle private_nh(private_nh_);
  private_nh.param("frame_id", m_worldFrameId, m_worldFrameId);
  private_nh.param("base_frame_id", m_baseFrameId, m_baseFrameId);

  private_nh.param("pointcloud_min_x", m_pointcloudMinX,m_pointcloudMinX);
  private_nh.param("pointcloud_max_x", m_pointcloudMaxX,m_pointcloudMaxX);
  private_nh.param("pointcloud_min_y", m_pointcloudMinY,m_pointcloudMinY);
  private_nh.param("pointcloud_max_y", m_pointcloudMaxY,m_pointcloudMaxY);
  private_nh.param("pointcloud_min_z", m_pointcloudMinZ,m_pointcloudMinZ);
  private_nh.param("pointcloud_max_z", m_pointcloudMaxZ,m_pointcloudMaxZ);
  private_nh.param("min_x_size", m_minSizeX,m_minSizeX);
  private_nh.param("min_y_size", m_minSizeY,m_minSizeY);

  private_nh.param("sensor_model/max_range", m_maxRange, m_maxRange);

  private_nh.param("resolution", m_res, m_res);
  private_nh.param("sensor_model/hit", m_probHit, m_probHit);
  private_nh.param("sensor_model/miss", m_probMiss, m_probMiss);
  private_nh.param("sensor_model/min", m_thresMin, m_thresMin);
  private_nh.param("sensor_model/max", m_thresMax, m_thresMax);

  // initialize octomap object & params
  m_octree = new OcTree(m_res);
  m_octree->setProbHit(m_probHit);
  m_octree->setProbMiss(m_probMiss);
  m_octree->setClampingThresMin(m_thresMin);
  m_octree->setClampingThresMax(m_thresMax);
  m_treeDepth = m_octree->getTreeDepth();
  m_maxTreeDepth = m_treeDepth;

  m_markerPub = m_nh.advertise<visualization_msgs::MarkerArray>("occupied_cells_vis_array", 1);
  m_binaryMapPub = m_nh.advertise<OctomapBinary>("octomap_binary", 1);
  m_fullMapPub = m_nh.advertise<OctomapBinary>("octomap_full", 1);

  m_pointCloudSub = new message_filters::Subscriber<sensor_msgs::PointCloud2> (m_nh, "cloud_in", 5);
  m_tfPointCloudSub = new tf::MessageFilter<sensor_msgs::PointCloud2> (*m_pointCloudSub, m_tfListener, m_worldFrameId, 5);
  m_tfPointCloudSub->registerCallback(boost::bind(&OctomapGen::insertCloudCallback, this, _1));

}

OctomapGen::~OctomapGen(){
  if (m_tfPointCloudSub){
    delete m_tfPointCloudSub;
    m_tfPointCloudSub = NULL;
  }

  if (m_pointCloudSub){
    delete m_pointCloudSub;
    m_pointCloudSub = NULL;
  }


  if (m_octree){
    delete m_octree;
    m_octree = NULL;
  }

}

void OctomapGen::insertCloudCallback(const sensor_msgs::PointCloud2::ConstPtr& cloud)
{
  ros::WallTime startTime = ros::WallTime::now();

  //
  // ground filtering in base frame
  //
  PCLPointCloud pc; // input cloud for filtering and ground-detection
  pcl::fromROSMsg(*cloud, pc);

  tf::StampedTransform sensorToWorldTf;
  try 
  {
    m_tfListener.lookupTransform(m_worldFrameId, cloud->header.frame_id, cloud->header.stamp, sensorToWorldTf);
  } 
  catch(tf::TransformException& ex)
  {
    ROS_ERROR_STREAM( "Transform error of sensor data: " << ex.what() << ", quitting callback");
    return;
  }

  Eigen::Matrix4f sensorToWorld;
  pcl_ros::transformAsMatrix(sensorToWorldTf, sensorToWorld);

  // set up filter for height range, also removes NANs:
  pcl::PassThrough<pcl::PointXYZ> passx;
  passx.setFilterFieldName("x");
  passx.setFilterLimits(m_pointcloudMinX, m_pointcloudMaxX);

  // set up filter for height range, also removes NANs:
  pcl::PassThrough<pcl::PointXYZ> passy;
  passy.setFilterFieldName("y");
  passy.setFilterLimits(m_pointcloudMinY, m_pointcloudMaxY);

  // set up filter for height range, also removes NANs:
  pcl::PassThrough<pcl::PointXYZ> passz;
  passz.setFilterFieldName("z");
  passz.setFilterLimits(m_pointcloudMinZ, m_pointcloudMaxZ);

  // directly transform to map frame:
  pcl::transformPointCloud(pc, pc, sensorToWorld);

  // just filter x range:
  passx.setInputCloud(pc.makeShared());
  passx.filter(pc);

  // just filter y range:
  passy.setInputCloud(pc.makeShared());
  passy.filter(pc);

  // just filter z range:
  passz.setInputCloud(pc.makeShared());
  passz.filter(pc);

  insertScan(sensorToWorldTf.getOrigin(), pc);

  double total_elapsed = (ros::WallTime::now() - startTime).toSec();
  ROS_DEBUG("Pointcloud insertion in OctomapGen done (%f sec)", total_elapsed);

//  if (m_markerPub.getNumSubscribers() > 0)
//  {
    // init markers:
    visualization_msgs::MarkerArray occupiedNodesVis;
    // each array stores all cubes of a different size, one for each depth level:
    occupiedNodesVis.markers.resize(m_treeDepth+1);
  
    // now, traverse all leafs in the tree:
    for (OcTree::iterator it = m_octree->begin(m_maxTreeDepth),
        end = m_octree->end(); it != end; ++it)
    {
      double x = it.getX();
      double y = it.getY();
      double z = it.getZ();
      
      if ((x > m_pointcloudMinX && x < m_pointcloudMaxX)
         && (y > m_pointcloudMinY && y < m_pointcloudMaxY)
         && (z > m_pointcloudMinZ && z < m_pointcloudMaxZ))
      {
    
        if (m_octree->isNodeOccupied(*it))
        {
          unsigned idx = it.getDepth();
          assert(idx < occupiedNodesVis.markers.size());

          geometry_msgs::Point cubeCenter;
          cubeCenter.x = x;
          cubeCenter.y = y;
          cubeCenter.z = z;

          occupiedNodesVis.markers[idx].points.push_back(cubeCenter);
      //    if (m_useHeightMap)
      //    {
      //      double minX, minY, minZ, maxX, maxY, maxZ;
      //      m_octree->getMetricMin(minX, minY, minZ);
      //      m_octree->getMetricMax(maxX, maxY, maxZ);

      //      double h = (1.0 - std::min(std::max((cubeCenter.z-minZ)/ (maxZ - minZ), 0.0), 1.0)) *m_colorFactor;
      //      occupiedNodesVis.markers[idx].colors.push_back(heightMapColor(h));
      //    }
         }
      }
    }

    // finish MarkerArray:
    for (unsigned i= 0; i < occupiedNodesVis.markers.size(); ++i)
    {
      double size = m_octree->getNodeSize(i);

      std_msgs::ColorRGBA color;
      color.r = 0.0;
      color.g = 1.0;
      color.b = 0.0;
      color.a = 1.0;

      occupiedNodesVis.markers[i].header.frame_id = m_worldFrameId;
      occupiedNodesVis.markers[i].header.stamp = cloud->header.stamp;
      occupiedNodesVis.markers[i].ns = "map";
      occupiedNodesVis.markers[i].id = i;
      occupiedNodesVis.markers[i].type = visualization_msgs::Marker::CUBE_LIST;
      occupiedNodesVis.markers[i].scale.x = size;
      occupiedNodesVis.markers[i].scale.y = size;
      occupiedNodesVis.markers[i].scale.z = size;
      occupiedNodesVis.markers[i].color = color;


      if (occupiedNodesVis.markers[i].points.size() > 0)
      {
        occupiedNodesVis.markers[i].action = visualization_msgs::Marker::ADD;
      }
      else
      {
        occupiedNodesVis.markers[i].action = visualization_msgs::Marker::DELETE;
      }
    }

    m_markerPub.publish(occupiedNodesVis);
//  }

  if (m_binaryMapPub.getNumSubscribers() > 0)
  {
    publishBinaryOctoMap(cloud->header.stamp);
  }

  if (m_fullMapPub.getNumSubscribers() > 0)
  {
    publishFullOctoMap(cloud->header.stamp);
  }

}

void OctomapGen::insertScan(const tf::Point& sensorOriginTf, const PCLPointCloud& pc)
{
  point3d sensorOrigin = pointTfToOctomap(sensorOriginTf);

  if (!m_octree->genKey(sensorOrigin, m_updateBBXMin)
      || !m_octree->genKey(sensorOrigin, m_updateBBXMax))
  {
    ROS_ERROR_STREAM("Could not generate Key for origin "<<sensorOrigin);
  }

  // all other points: free on ray, occupied on endpoint:
  KeySet free_cells, occupied_cells;
  for (PCLPointCloud::const_iterator it = pc.begin(); it != pc.end(); ++it){
    point3d point(it->x, it->y, it->z);
    // maxrange check
    if ((m_maxRange < 0.0) || ((point - sensorOrigin).norm() <= m_maxRange) ) {

      // free cells
      if (m_octree->computeRayKeys(sensorOrigin, point, m_keyRay)){
        free_cells.insert(m_keyRay.begin(), m_keyRay.end());
      }
      // occupied endpoint
      OcTreeKey key;
      if (m_octree->genKey(point, key)){
        occupied_cells.insert(key);

        updateMinKey(key, m_updateBBXMin);
        updateMaxKey(key, m_updateBBXMax);
      }
    } 
    else 
    { // ray longer than maxrange:;
      point3d new_end = sensorOrigin + (point - sensorOrigin).normalized() * m_maxRange;
      if (m_octree->computeRayKeys(sensorOrigin, new_end, m_keyRay)){
        free_cells.insert(m_keyRay.begin(), m_keyRay.end());

        octomap::OcTreeKey endKey;
        if (m_octree->genKey(new_end, endKey))
        {
          updateMinKey(endKey, m_updateBBXMin);
          updateMaxKey(endKey, m_updateBBXMax);
        }
        else
        {
          ROS_ERROR_STREAM("Could not generate Key for endpoint "<<new_end);
        }

      }
    }
  }

  // mark free cells only if not seen occupied in this cloud
  for(KeySet::iterator it = free_cells.begin(), end=free_cells.end(); it!= end; ++it)
  {
    if (occupied_cells.find(*it) == occupied_cells.end())
    {
      m_octree->updateNode(*it, false);
    }
  }

  // now mark all occupied cells:
  for (KeySet::iterator it = occupied_cells.begin(), end=free_cells.end(); it!= end; it++)
  {
    m_octree->updateNode(*it, true);
  }

  // TODO: eval lazy+updateInner vs. proper insertion
  // non-lazy by default (updateInnerOccupancy() too slow for large maps)
  //m_octree->updateInnerOccupancy();
  octomap::point3d minPt, maxPt;
  ROS_DEBUG_STREAM("Bounding box keys (before): " << m_updateBBXMin[0] << " " <<m_updateBBXMin[1] << " " << m_updateBBXMin[2] << " / " <<m_updateBBXMax[0] << " "<<m_updateBBXMax[1] << " "<< m_updateBBXMax[2]);

  // TODO: snap max / min keys to larger voxels by m_maxTreeDepth
//   if (m_maxTreeDepth < 16)
//   {
//      OcTreeKey tmpMin = getIndexKey(m_updateBBXMin, m_maxTreeDepth); // this should give us the first key at depth m_maxTreeDepth that is smaller or equal to m_updateBBXMin (i.e. lower left in 2D grid coordinates)
//      OcTreeKey tmpMax = getIndexKey(m_updateBBXMax, m_maxTreeDepth); // see above, now add something to find upper right
//      tmpMax[0]+= m_octree->getNodeSize( m_maxTreeDepth ) - 1;
//      tmpMax[1]+= m_octree->getNodeSize( m_maxTreeDepth ) - 1;
//      tmpMax[2]+= m_octree->getNodeSize( m_maxTreeDepth ) - 1;
//      m_updateBBXMin = tmpMin;
//      m_updateBBXMax = tmpMax;
//   }

  // TODO: we could also limit the bbx to be within the map bounds here (see publishing check)
  m_octree->genCoords(m_updateBBXMin, m_octree->getTreeDepth(), minPt);
  m_octree->genCoords(m_updateBBXMax, m_octree->getTreeDepth(), maxPt);
  ROS_DEBUG_STREAM("Updated area bounding box: "<< minPt << " - "<<maxPt);
  ROS_DEBUG_STREAM("Bounding box keys (after): " << m_updateBBXMin[0] << " " <<m_updateBBXMin[1] << " " << m_updateBBXMin[2] << " / " <<m_updateBBXMax[0] << " "<<m_updateBBXMax[1] << " "<< m_updateBBXMax[2]);

//  if (m_compressMap)
//    m_octree->prune();


}


void OctomapGen::publishBinaryOctoMap(const ros::Time& rostime) const{

  OctomapBinary map;
  map.header.frame_id = m_worldFrameId;
  map.header.stamp = rostime;

  if (octomap_msgs::binaryMapToMsgData(*m_octree, map.data))
    m_binaryMapPub.publish(map);
  else
    ROS_ERROR("Error serializing OctoMap");
}

void OctomapGen::publishFullOctoMap(const ros::Time& rostime) const{

  OctomapBinary map;
  map.header.frame_id = m_worldFrameId;
  map.header.stamp = rostime;

  if (octomap_msgs::fullMapToMsgData(*m_octree, map.data))
    m_fullMapPub.publish(map);
  else
    ROS_ERROR("Error serializing OctoMap");

}

// ###### just for visualisation purpuses #######


}
